// rlezip.c - RLE 压缩程序
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>     // 用于计时
#include <windows.h>  // 用于 FindFirstFile 等以及 GetFileAttributes
#include <sys/stat.h> // 用于 _stat64 获取文件大小
#include "rle_common.h"

// 前向声明
static long long total_original_size = 0;
static long long total_compressed_size = 0;

// 获取文件大小的函数
long long get_file_size(const char *filename) {
    struct __stat64 stat_buf;
    if (_stat64(filename, &stat_buf) == 0) {
        return stat_buf.st_size;
    }
    fprintf(stderr, "Error getting file size for %s: ", filename);
    perror(NULL);
    return -1;
}

// 压缩单个文件并将其添加到存档的函数
// disk_file_path: 磁盘上文件的完整路径 (例如, "C:/Users/X/file.txt")
// archive_relative_path: 要存储在头文件中的路径 (例如, "file.txt" 或 "somedir/file.txt")
int add_file_to_archive(FILE *archive_fp, const char *disk_file_path, const char *archive_relative_path) {
    RzpEntryHeader header;
    FILE *input_fp = NULL;
    long file_pos_before_header, file_pos_after_data;

    input_fp = fopen(disk_file_path, "rb");
    if (!input_fp) {
        fprintf(stderr, "Error: Cannot open input file %s\n", disk_file_path);
        return -1;
    }

    memset(&header, 0, sizeof(RzpEntryHeader));
    // 提取文件名用于 header.name
    const char *filename_ptr = strrchr(archive_relative_path, '/');
    if (filename_ptr) {
        filename_ptr++; // 移动到斜杠之后
    } else {
        filename_ptr = archive_relative_path; // 没有斜杠，是顶级文件
    }
    strncpy(header.name, filename_ptr, MAX_FILENAME_LEN - 1);
    header.name[MAX_FILENAME_LEN - 1] = '\0';

    strncpy(header.relative_path, archive_relative_path, sizeof(header.relative_path) - 1);
    header.relative_path[sizeof(header.relative_path) - 1] = '\0';

    header.type = RLE_FILE_TYPE;
    header.original_size = get_file_size(disk_file_path);
    if (header.original_size == -1) {
        fclose(input_fp);
        return -1;
    }
    header.compressed_size = 0; // 占位符

    file_pos_before_header = ftell(archive_fp);
    if (file_pos_before_header == -1L) { perror("ftell before header write failed"); fclose(input_fp); return -1; }
    write_entry_header(archive_fp, &header); // 写入占位符头
    
    if (rle_compress(input_fp, archive_fp, &header.compressed_size) != 0) {
        fprintf(stderr, "Error: Compression failed for file %s\n", disk_file_path);
        fclose(input_fp);
        // 注意: 存档文件此时可能处于不一致状态。
        // 正确的错误处理可能包括尝试将存档文件截断回 file_pos_before_header。
        // 目前，我们将继续，但出错时存档可能会损坏。
        return -1;
    }
    file_pos_after_data = ftell(archive_fp);
    if (file_pos_after_data == -1L) { perror("ftell after data write failed"); fclose(input_fp); return -1; }

    // 使用正确的 compressed_size 更新头文件
    if (fseek(archive_fp, file_pos_before_header, SEEK_SET) != 0) { perror("fseek to update header failed"); fclose(input_fp); return -1; }
    write_entry_header(archive_fp, &header); // 写入实际的头文件
    if (fseek(archive_fp, file_pos_after_data, SEEK_SET) != 0) { perror("fseek to end of data failed"); fclose(input_fp); return -1; }

    fclose(input_fp);
    printf("Compressed file: %s (Original: %I64d, Compressed: %I64d, Archive Path: %s)\n", 
           disk_file_path, header.original_size, header.compressed_size, archive_relative_path);
    total_original_size += header.original_size;
    total_compressed_size += header.compressed_size; 
    total_compressed_size += sizeof(RzpEntryHeader);
    return 0;
}

// 将目录（递归地）添加到存档的函数
// disk_path_of_dir_to_compress: 磁盘上要压缩目录的完整路径 (例如, "C:/data/mydir")
// archive_relative_path_for_this_dir: 此目录在存档中的相对路径 (例如, "mydir/" 或 "mydir/subdir/")
int add_directory_to_archive(FILE *archive_fp, const char *disk_path_of_dir_to_compress, const char *archive_relative_path_for_this_dir) {
    RzpEntryHeader dir_header;
    WIN32_FIND_DATAA find_data;
    HANDLE hFind = INVALID_HANDLE_VALUE;
    char search_pattern[MAX_PATH];
    char child_disk_path[MAX_PATH];
    char child_archive_relative_path[MAX_FILENAME_LEN * 2]; // 头文件中相对路径的最大长度

    // 1. 写入当前目录本身的头文件
    memset(&dir_header, 0, sizeof(RzpEntryHeader));
    dir_header.type = RLE_DIRECTORY_TYPE;
    dir_header.original_size = 0;
    dir_header.compressed_size = 0;

    strncpy(dir_header.relative_path, archive_relative_path_for_this_dir, sizeof(dir_header.relative_path) - 1);
    dir_header.relative_path[sizeof(dir_header.relative_path) - 1] = '\0';

    // 提取目录名用于 header.name (例如, 从 "mydir/" 中提取 "mydir", 或从 "somedir/subdir1/" 中提取 "subdir1")
    char temp_name_path[MAX_FILENAME_LEN * 2];
    strncpy(temp_name_path, archive_relative_path_for_this_dir, sizeof(temp_name_path) - 1);
    temp_name_path[sizeof(temp_name_path) - 1] = '\0';
    int len = strlen(temp_name_path);
    if (len > 0 && temp_name_path[len - 1] == '/') {
        temp_name_path[len - 1] = '\0'; // 删除末尾的斜杠以提取名称
    }
    const char *last_slash = strrchr(temp_name_path, '/');
    if (last_slash) {
        strncpy(dir_header.name, last_slash + 1, MAX_FILENAME_LEN - 1);
    } else {
        strncpy(dir_header.name, temp_name_path, MAX_FILENAME_LEN - 1);
    }
    dir_header.name[MAX_FILENAME_LEN - 1] = '\0';

    write_entry_header(archive_fp, &dir_header);
    total_compressed_size += sizeof(RzpEntryHeader);
    printf("Added directory entry: %s (Archive Path: %s)\n", dir_header.name, dir_header.relative_path);

    // 2. 遍历磁盘上目录的内容
    snprintf(search_pattern, MAX_PATH, "%s\\*", disk_path_of_dir_to_compress); // Windows 磁盘路径使用 \

    hFind = FindFirstFileA(search_pattern, &find_data);
    if (hFind == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "Error: FindFirstFileA failed for %s (Error code: %lu)\n", search_pattern, GetLastError());
        return -1; 
    }

    do {
        const char *found_name = find_data.cFileName;
        if (strcmp(found_name, ".") == 0 || strcmp(found_name, "..") == 0) {
            continue; // 跳过当前目录和父目录条目
        }

        snprintf(child_disk_path, MAX_PATH, "%s\\%s", disk_path_of_dir_to_compress, found_name);

        if (find_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
            snprintf(child_archive_relative_path, sizeof(child_archive_relative_path), "%s%s/", archive_relative_path_for_this_dir, found_name);
            if (add_directory_to_archive(archive_fp, child_disk_path, child_archive_relative_path) != 0) {
                // 递归调用出错，传播错误
                FindClose(hFind);
                return -1;
            }
        } else { // 是文件
            snprintf(child_archive_relative_path, sizeof(child_archive_relative_path), "%s%s", archive_relative_path_for_this_dir, found_name);
            if (add_file_to_archive(archive_fp, child_disk_path, child_archive_relative_path) != 0) {
                // 压缩文件出错，传播错误
                FindClose(hFind);
                return -1;
            }
        }
    } while (FindNextFileA(hFind, &find_data) != 0);

    DWORD dwError = GetLastError();
    FindClose(hFind);
    if (dwError != ERROR_NO_MORE_FILES) {
        fprintf(stderr, "Error: FindNextFileA failed with error code %lu\n", dwError);
        return -1;
    }

    return 0;
}


int main(int argc, char *argv[]) {
    clock_t start_time, end_time;
    double cpu_time_used;

    if (argc < 3) {
        fprintf(stderr, "Usage: rlezip <input_file(s)/dir> ... <output_file.rzp>\n");
        fprintf(stderr, "Examples:\n");
        fprintf(stderr, "  rlezip file.txt archive.rzp\n");
        fprintf(stderr, "  rlezip file1.txt file2.txt archive.rzp\n");
        fprintf(stderr, "  rlezip * archive.rzp (压缩当前目录中的所有文件，不包括子目录)\n");
        fprintf(stderr, "  rlezip directory archive.rzp (压缩目录)\n");
        fprintf(stderr, "  rlezip C:\\path\\to\\directory archive.rzp (通过绝对路径压缩目录)\n");
        return 1;
    }

    start_time = clock();

    const char *output_filename = argv[argc - 1];
    FILE *archive_fp = fopen(output_filename, "wb");
    if (!archive_fp) {
        perror("Error: Cannot open output archive file");
        return 1;
    }

    printf("rlezip program started. Output archive: %s\n", output_filename);

    for (int i = 1; i < argc - 1; ++i) {
        const char *input_path_arg = argv[i];
        printf("Processing argument: %s\n", input_path_arg);

        if (strcmp(input_path_arg, "*") == 0) {
            WIN32_FIND_DATAA find_data_wildcard;
            HANDLE hFind_wildcard = INVALID_HANDLE_VALUE;
            char current_dir_search_pattern[MAX_PATH];
            // char file_to_add_path[MAX_PATH]; // Unused variable

            // Get current directory for FindFirstFileA
            // Using ".\\*" to search in current directory
            snprintf(current_dir_search_pattern, MAX_PATH, ".\\*");

            hFind_wildcard = FindFirstFileA(current_dir_search_pattern, &find_data_wildcard);
            if (hFind_wildcard == INVALID_HANDLE_VALUE) {
                fprintf(stderr, "Error: FindFirstFileA failed for wildcard '*' in current directory (Error code: %lu)\n", GetLastError());
                // continue; // Or handle error as appropriate
            } else {
                printf("Processing wildcard '*' in current directory...\n");
                do {
                    const char *found_filename = find_data_wildcard.cFileName;
                    if (strcmp(found_filename, ".") == 0 || strcmp(found_filename, "..") == 0) {
                        continue; // Skip . and ..
                    }
                    // Check if it's a directory. If so, skip it as per requirement "不考虑子目录"
                    if (find_data_wildcard.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
                        printf("Skipping directory (due to '*'): %s\n", found_filename);
                        continue;
                    }
                    // Construct full path for file to add (it's in current dir)
                    // For add_file_to_archive, disk_file_path is just the filename (relative to CWD)
                    // and archive_relative_path is also just the filename.
                    // No, disk_file_path needs to be relative to CWD if CWD is not where the exe is.
                    // However, fopen and _stat64 usually resolve relative paths from CWD.
                    // Let's assume filename is sufficient if CWD is where files are.
                    
                    // The problem states "将当前目录下的所有文件压缩", so found_filename is the path on disk
                    // and also the path in the archive root.
                    if (add_file_to_archive(archive_fp, found_filename, found_filename) != 0) {
                        fprintf(stderr, "Error processing file (from '*'): %s\n", found_filename);
                    }

                } while (FindNextFileA(hFind_wildcard, &find_data_wildcard) != 0);
                
                DWORD dwErrorWildcard = GetLastError();
                FindClose(hFind_wildcard);
                if (dwErrorWildcard != ERROR_NO_MORE_FILES) {
                    fprintf(stderr, "Error: FindNextFileA (wildcard) failed with error code %lu\n", dwErrorWildcard);
                }
            }
        } else {
            DWORD attributes = GetFileAttributesA(input_path_arg);
            if (attributes == INVALID_FILE_ATTRIBUTES) {
                fprintf(stderr, "Error: Cannot access input path %s or path does not exist. (Code: %lu)\n", input_path_arg, GetLastError());
                continue; 
            }

            char path_basename[MAX_FILENAME_LEN];
            char temp_input_path_for_basename[MAX_PATH];
            strncpy(temp_input_path_for_basename, input_path_arg, MAX_PATH -1);
            temp_input_path_for_basename[MAX_PATH-1] = '\0';
            
            // Remove trailing slash if present, to get the actual directory name for basename
            int path_len = strlen(temp_input_path_for_basename);
            if (path_len > 0 && (temp_input_path_for_basename[path_len-1] == '\\' || temp_input_path_for_basename[path_len-1] == '/')) {
                if (path_len > 1 && temp_input_path_for_basename[path_len-2] != ':') { // Avoid removing slash from "C:\"
                     temp_input_path_for_basename[path_len-1] = '\0';
                }
            }

            const char *name_ptr = strrchr(temp_input_path_for_basename, '\\');
            if (!name_ptr) name_ptr = strrchr(temp_input_path_for_basename, '/');
            if (name_ptr) name_ptr++; 
            else name_ptr = temp_input_path_for_basename; 
            strncpy(path_basename, name_ptr, MAX_FILENAME_LEN - 1);
            path_basename[MAX_FILENAME_LEN - 1] = '\0';

            if (attributes & FILE_ATTRIBUTE_DIRECTORY) {
                char archive_initial_relative_path[MAX_FILENAME_LEN * 2];
                // For a directory 'somedir', its archive relative path is "somedir/"
                snprintf(archive_initial_relative_path, sizeof(archive_initial_relative_path), "%s/", path_basename);
                if (add_directory_to_archive(archive_fp, input_path_arg, archive_initial_relative_path) != 0) {
                    fprintf(stderr, "Error processing directory %s\n", input_path_arg);
                    // Decide on error handling: continue with next arg or abort?
                    // For now, continue.
                }
            } else { // It's a file
                // For a single file argument, its archive path is just its basename.
                if (add_file_to_archive(archive_fp, input_path_arg, path_basename) != 0) {
                     fprintf(stderr, "Error processing file %s\n", input_path_arg);
                }
            }
        }
    }

    fclose(archive_fp);
    end_time = clock();
    cpu_time_used = ((double) (end_time - start_time)) / CLOCKS_PER_SEC;

    printf("----------------------------------------\n");
    printf("Compression Summary:\n");
    printf("Total original size: %I64d bytes\n", total_original_size);
    printf("Total compressed size (including headers): %I64d bytes\n", total_compressed_size);
    if (total_original_size > 0 && total_compressed_size > 0) {
        // Compression ratio: (original_size / compressed_size) : 1
        // Or, space saving: (1 - compressed_size / original_size) * 100%
        // Or, compressed is X% of original: (compressed_size / original_size) * 100%
        // The problem asks for "压缩比", typically (compressed / original) or (original / compressed).
        // Let's use (compressed / original) * 100%.
        double ratio_percent = (double)total_compressed_size / total_original_size * 100.0;
        printf("Archive is %.2f%% of original size.\n", ratio_percent);
    } else if (total_original_size == 0 && total_compressed_size > 0) {
        printf("Compression ratio: N/A (original size is zero, archive has metadata)\n");
    } else if (total_original_size > 0 && total_compressed_size == 0) {
        // This case should ideally not happen if headers are counted.
        printf("Compression ratio: N/A (compressed to zero size? Check logic.)\n"); 
    } else {
        printf("Compression ratio: N/A (no data processed or empty archive)\n");
    }
    printf("Time taken: %.4f seconds\n", cpu_time_used);
    printf("rlezip program finished.\n");

    return 0;
}